#include "TcpChannel.h"
#include "NetworkMessage.h"
#include <algorithm>
#include <chrono>
#include <functional>

namespace SeanNetwork {

TcpChannel::TcpChannel(asio::io_context& ioContext)
    : BaseNetworkChannel(ioContext),
      m_Socket(ioContext),
      m_HeaderSize(MESSAGE_HEADER_SIZE),
      m_KeepAliveInterval(15000)        // 默认15秒心跳间隔
{
    m_ReadBuffer.resize(4096);  // 初始接收缓冲区大小
    
    // 设置消息确认回调
    m_MessagePool->SetAckCallback([this](uint32_t messageId, bool success) {
        HandleMessageAck(messageId, success);
    });
}

TcpChannel::~TcpChannel() {
    DoDisconnect();
}

bool TcpChannel::Initialize() {
    return BaseNetworkChannel::Initialize();
}

bool TcpChannel::Connect(const std::string& host, uint16_t port) {
    return BaseNetworkChannel::Connect(host, port);
}

void TcpChannel::Disconnect() {
    BaseNetworkChannel::Disconnect();
}

void TcpChannel::Update() {
    BaseNetworkChannel::Update();
}

void TcpChannel::DoSendMessage(std::shared_ptr<NetworkMessage> message) {
    // 在strand中发送，确保线程安全
    asio::post(m_Strand, [this, message]() {
        if (m_Socket.is_open() && m_State == ConnectionState::Connected) {
            SendNextMessage();
        }
    });
}

bool TcpChannel::DoInitialize() {
    /* 功能：初始化TCP通道
       返回：初始化是否成功 */
    try {
        if (m_KeepAliveTimer) {
            asio::error_code ec;
            m_KeepAliveTimer->cancel(ec);
        }
        
        m_KeepAliveTimer = std::make_shared<asio::steady_timer>(m_IoContext);
        
        if (m_ReconnectTimer) {
            asio::error_code ec;
            m_ReconnectTimer->cancel(ec);
        }
        
        m_ReconnectTimer = std::make_shared<asio::steady_timer>(m_IoContext);
        
        return true;
    }
    catch (const std::exception& e) {
        UpdateConnectionState(ConnectionState::Failed, std::string("Initialize error: ") + e.what());
        return false;
    }
}

bool TcpChannel::DoConnect(const std::string& host, uint16_t port) {
    /* 功能：连接到远程主机
       参数：host-主机地址，port-端口号
       返回：连接操作是否成功启动 */
    if (m_State == ConnectionState::Connected || m_State == ConnectionState::Connecting) {
        return false;
    }
    
    m_Host = host;
    m_Port = port;
    m_ReconnectAttempts = 0;
    
    UpdateConnectionState(ConnectionState::Connecting);
    
    try {
        // 解析地址
        asio::ip::tcp::resolver resolver(m_IoContext);
        auto endpoints = resolver.resolve(host, std::to_string(port));
        
        // 关闭旧连接
        asio::error_code ec;
        if (m_Socket.is_open()) {
            m_Socket.close(ec);
        }
        
        // 异步连接
        asio::async_connect(m_Socket, endpoints,
            asio::bind_executor(m_Strand, [this](const asio::error_code& ec, const asio::ip::tcp::endpoint& endpoint) {
                if (!ec) {
                    m_Endpoint = endpoint;
                    
                    // 设置套接字选项
                    asio::error_code optEc;
                    m_Socket.set_option(asio::ip::tcp::no_delay(true), optEc);
                    
                    UpdateConnectionState(ConnectionState::Connected);
                    
                    // 启动接收和心跳定时器
                    StartReceive();
                    SendKeepAlive();
                }
                else {
                    std::string errMsg = "Connect error: " + ec.message();
                    UpdateConnectionState(ConnectionState::Failed, errMsg);
                    
                    // 启动重连
                    StartReconnect();
                }
            }));
        
        return true;
    }
    catch (const std::exception& e) {
        UpdateConnectionState(ConnectionState::Failed, std::string("Connect exception: ") + e.what());
        StartReconnect();
        return false;
    }
}

void TcpChannel::DoDisconnect() {
    /* 功能：断开连接 */
    if (m_State == ConnectionState::Disconnected) {
        return;
    }
    
    // 取消定时器
    if (m_KeepAliveTimer) {
        asio::error_code ec;
        m_KeepAliveTimer->cancel(ec);
    }
    
    if (m_ReconnectTimer) {
        asio::error_code ec;
        m_ReconnectTimer->cancel(ec);
    }
    
    // 关闭套接字
    asio::error_code ec;
    if (m_Socket.is_open()) {
        m_Socket.shutdown(asio::ip::tcp::socket::shutdown_both, ec);
        m_Socket.close(ec);
    }
    
    // 清空发送队列
    {
        std::lock_guard<std::mutex> lock(m_SendQueueMutex);
        m_SendQueue.clear();
    }
    
    // 清空消息池
    if (m_MessagePool) {
        m_MessagePool->Clear();
    }
    
    UpdateConnectionState(ConnectionState::Disconnected);
}

bool TcpChannel::SendMsg(std::shared_ptr<NetworkMessage> message) {
    /* 功能：发送消息
       参数：message-要发送的消息
       返回：发送操作是否成功启动 */
    if (!message || m_State != ConnectionState::Connected) {
        return false;
    }
    
    message->UpdateTimestamp();
    message->UpdateLastSendTime();
    
    // 添加到消息池进行可靠性管理
    if (m_MessagePool->AddMessage(message, ChannelType::TCP)) {
        // 加入发送队列
        {
            std::lock_guard<std::mutex> lock(m_SendQueueMutex);
            m_SendQueue.push_back(message);
        }
        
        // 在strand中发送，确保线程安全
        asio::post(m_Strand, [this]() {
            if (m_Socket.is_open() && m_State == ConnectionState::Connected) {
                if (m_SendQueue.size() == 1) {  // 只有一个待发送时，启动发送流程
                    SendNextMessage();
                }
            }
        });
        
        return true;
    }
    
    return false;
}

void TcpChannel::SetMessageCallback(MessageReceiveCallback callback) {
    m_MessageCallback = callback;
}

void TcpChannel::SetConnectionCallback(ConnectionStateCallback callback) {
    m_ConnectionCallback = callback;
}

ConnectionState TcpChannel::GetConnectionState() const {
    return m_State;
}

NetworkStats TcpChannel::GetNetworkStats() const {
    std::lock_guard<std::mutex> lock(m_StatsMutex);
    return m_Stats;
}

ChannelType TcpChannel::GetChannelType() const {
    return ChannelType::TCP;
}

uint32_t TcpChannel::GetLatency() const {
    std::lock_guard<std::mutex> lock(m_StatsMutex);
    return m_Stats.LatestRtt;
}

void TcpChannel::DoUpdate() {
    /* 功能：更新通道状态
       处理：检查重传和超时消息 */
    if (m_State != ConnectionState::Connected) {
        return;
    }
    
    // 检查消息超时
    auto timeoutMessages = m_MessagePool->CheckTimeoutMessages();
    
    // 重新发送超时消息
    for (auto& message : timeoutMessages) {
        std::lock_guard<std::mutex> lock(m_SendQueueMutex);
        m_SendQueue.push_front(message);  // 超时消息优先发送
    }
    
    // 如果有待发送消息，确保发送流程启动
    if (!m_SendQueue.empty()) {
        asio::post(m_Strand, [this]() {
            if (m_Socket.is_open() && m_State == ConnectionState::Connected) {
                if (!m_SendQueue.empty()) {
                    SendNextMessage();
                }
            }
        });
    }
}

void TcpChannel::SetKeepAliveInterval(uint32_t intervalMs) {
    m_KeepAliveInterval = intervalMs;
    
    // 更新现有定时器
    if (m_State == ConnectionState::Connected && m_KeepAliveTimer) {
        m_KeepAliveTimer->expires_after(std::chrono::milliseconds(m_KeepAliveInterval));
        m_KeepAliveTimer->async_wait(asio::bind_executor(m_Strand, [this](const asio::error_code& ec) {
            if (!ec) {
                SendKeepAlive();
            }
        }));
    }
}

void TcpChannel::SetReconnectStrategy(uint32_t initialDelayMs, uint32_t maxDelayMs, uint32_t maxAttempts) {
    if (initialDelayMs == 0 || maxDelayMs == 0) {
        return;
    }
    
    m_ReconnectInitialDelay = initialDelayMs;
    m_ReconnectMaxDelay = maxDelayMs;
    m_ReconnectMaxAttempts = maxAttempts;
}

void TcpChannel::StartReceive() {
    /* 功能：启动接收流程 */
    if (!m_Socket.is_open() || m_State != ConnectionState::Connected) {
        return;
    }
    
    ReadHeader();
}

void TcpChannel::ReadHeader() {
    /* 功能：读取消息头
       描述：异步读取消息头，完成后解析消息长度并继续读取消息体 */
    
    // 至少要读取消息头
    if (m_ReadBuffer.size() < m_HeaderSize) {
        m_ReadBuffer.resize(m_HeaderSize);
    }
    
    // 用strand保证回调的线程安全
    asio::async_read(m_Socket,
        asio::buffer(m_ReadBuffer.data(), m_HeaderSize),
        asio::transfer_exactly(m_HeaderSize),
        asio::bind_executor(m_Strand,
            [this](asio::error_code ec, std::size_t bytes_transferred) {
                if (!ec && bytes_transferred == m_HeaderSize) {
                    // 解析消息头
                    MessageHeader header;
                    std::memcpy(&header, m_ReadBuffer.data(), m_HeaderSize);
                    
                    // 验证消息长度的合理性
                    const uint32_t MAX_MESSAGE_SIZE = 10 * 1024 * 1024; // 10MB上限
                    if (header.MessageLength < m_HeaderSize || header.MessageLength > MAX_MESSAGE_SIZE) {
                        // 消息长度不合理，可能是恶意攻击或数据损坏，断开连接
                        UpdateConnectionState(ConnectionState::Failed, "Invalid message length");
                        DoDisconnect();
                        return;
                    }
                    
                    // 继续读取消息体
                    ReadBody(header.MessageLength);
                } else {
                    // 读取失败，处理错误
                    std::string errorMsg = ec ? ec.message() : "Unexpected bytes transferred";
                    UpdateConnectionState(ConnectionState::Failed, "Header read error: " + errorMsg);
                    DoDisconnect();
                }
            }));
}

void TcpChannel::ReadBody(uint32_t messageLength) {
    /* 功能：读取消息体
       参数：messageLength-消息总长度（包括头部）
       描述：异步读取消息体，完成后处理完整消息 */
    
    if (messageLength <= m_HeaderSize) {
        // 没有消息体，直接开始读取下一条消息
        ReadHeader();
        return;
    }
    
    // 计算消息体大小并调整缓冲区
    uint32_t bodyLength = messageLength - m_HeaderSize;
    
    // 验证消息体长度的合理性
    const uint32_t MAX_BODY_SIZE = 10 * 1024 * 1024 - m_HeaderSize; // 10MB上限
    if (bodyLength > MAX_BODY_SIZE) {
        // 消息体过大，可能是恶意攻击，断开连接
        UpdateConnectionState(ConnectionState::Failed, "Message body too large");
        DoDisconnect();
        return;
    }
    
    try {
        // 确保缓冲区足够大
        if (m_ReadBuffer.size() < messageLength) {
            m_ReadBuffer.resize(messageLength);
        }
        
        // 异步读取消息体
        asio::async_read(m_Socket,
            asio::buffer(m_ReadBuffer.data() + m_HeaderSize, bodyLength),
            asio::transfer_exactly(bodyLength),
            asio::bind_executor(m_Strand,
                [this, messageLength](asio::error_code ec, std::size_t bytes_transferred) {
                    if (!ec) {
                        // 成功读取消息，进行处理
                        HandleCompleteMessage(messageLength);
                        
                        // 继续读取下一条消息
                        ReadHeader();
                    } else {
                        // 读取失败，处理错误
                        UpdateConnectionState(ConnectionState::Failed, "Body read error: " + ec.message());
                        DoDisconnect();
                    }
                }));
    }
    catch (const std::exception& e) {
        // 捕获所有可能的异常，包括内存分配失败
        UpdateConnectionState(ConnectionState::Failed, std::string("Exception in ReadBody: ") + e.what());
        DoDisconnect();
    }
}

// 新增方法：处理完整消息
void TcpChannel::HandleCompleteMessage(uint32_t messageLength) {
    try {
        // 使用零复制反序列化，避免不必要的内存复制
        auto message = NetworkMessage::DeserializeZeroCopy(m_ReadBuffer.data(), messageLength);
        
        if (!message) {
            // 反序列化失败，可能是无效消息，但我们继续读取
            BaseNetworkChannel::UpdateNetworkStats(0, messageLength, false, false);
            return;
        }
        
        // 更新统计信息
        BaseNetworkChannel::UpdateNetworkStats(0, messageLength, false, true);
        
        // 更新RTT（如果是ping包响应的话）
        if (message->HasFlag(MessageFlags::FLAG_KEEPALIVE) && message->GetHeader().MessageId == 0) {
            // 如果是心跳响应包，更新RTT
            uint64_t sendTime = message->GetHeader().Timestamp;
            uint64_t currentTime = GetCurrentTimeMillis();
            if (currentTime > sendTime) {
                BaseNetworkChannel::UpdateRtt(static_cast<uint32_t>(currentTime - sendTime));
            }
        }
        
        // 使用基类中的线程安全回调函数
        BaseNetworkChannel::InvokeMessageCallback(message);
    }
    catch (const std::exception& e) {
        // 捕获并记录异常，但继续处理下一条消息
        LOG_ERROR("HandleCompleteMessage异常: " + std::string(e.what()));
    }
}

void TcpChannel::SendNextMessage() {
    /* 功能：发送队列中的下一条消息 */
    std::shared_ptr<NetworkMessage> message = nullptr;
    
    // 首先尝试从无锁队列获取消息
    if (m_MessagePool->DequeueOutgoingMessage(message)) {
        // 使用无锁队列获取到消息，直接处理
    } else {
        // 如果无锁队列为空，则从优先级队列获取
        std::lock_guard<std::mutex> lock(m_SendQueueMutex);
        if (m_SendQueue.empty()) {
            return;
        }
        
        message = m_SendQueue.front();
        m_SendQueue.pop_front();
    }
    
    if (!message || !m_Socket.is_open() || m_State != ConnectionState::Connected) {
        return;
    }
    
    // 序列化消息并存储在共享指针中，确保数据在异步操作期间存活
    auto serializedData = message->Serialize();
    auto serializedDataPtr = std::make_shared<decltype(serializedData)>(std::move(serializedData));
    
    // 异步发送
    asio::async_write(m_Socket, asio::buffer(*serializedDataPtr),
        asio::bind_executor(m_Strand, [this, message, serializedDataPtr](const asio::error_code& ec, std::size_t bytes_transferred) {
            if (!ec) {
                // 发送成功，更新统计信息
                UpdateNetworkStats(bytes_transferred, 0, true, false);
                
                // 继续发送下一条消息
                std::shared_ptr<NetworkMessage> nextMessage = nullptr;
                if (m_MessagePool->DequeueOutgoingMessage(nextMessage) || !m_SendQueue.empty()) {
                    SendNextMessage();
                }
            }
            else {
                if (ec != asio::error::operation_aborted) {
                    // 发送失败，连接可能断开
                    std::string errMsg = "Send error: " + ec.message();
                    UpdateConnectionState(ConnectionState::Failed, errMsg);
                    
                    // 启动重连
                    StartReconnect();
                }
            }
        }));
}

void TcpChannel::HandleMessageAck(uint32_t messageId, bool success) {
    /* 功能：处理消息确认
       参数：messageId-消息ID，success-是否成功 */
    if (success) {
        // 消息成功送达
    }
    else {
        // 消息传输失败，超时或放弃
        std::lock_guard<std::mutex> lock(m_StatsMutex);
        m_Stats.PacketLoss++;
    }
}

void TcpChannel::SendKeepAlive() {
    /* 功能：发送心跳包
       处理：创建并发送心跳消息，设置下一次心跳定时器 */
    if (m_State != ConnectionState::Connected) {
        return;
    }
    
    // 创建心跳消息
    auto keepAliveMsg = std::make_shared<NetworkMessage>(0, MessagePriority::High);  // 0为心跳消息ID
    keepAliveMsg->SetBody(nullptr, 0);  // 空消息体
    
    // 发送心跳
    SendMsg(keepAliveMsg);
    
    // 设置下一次心跳
    if (m_KeepAliveTimer) {
        m_KeepAliveTimer->expires_after(std::chrono::milliseconds(m_KeepAliveInterval));
        m_KeepAliveTimer->async_wait(asio::bind_executor(m_Strand, [this](const asio::error_code& ec) {
            if (!ec) {
                SendKeepAlive();
            }
        }));
    }
}

void TcpChannel::StartReconnect() {
    /* 功能：启动重连流程
       处理：根据重连策略设置延迟后尝试重连 */
    if (m_State == ConnectionState::Reconnecting) {
        return;
    }
    
    if (m_ReconnectAttempts >= m_ReconnectMaxAttempts) {
        UpdateConnectionState(ConnectionState::Disconnected, "Maximum reconnect attempts reached");
        return;
    }
    
    UpdateConnectionState(ConnectionState::Reconnecting, "Connection lost, attempting to reconnect");
    
    m_ReconnectAttempts++;
    uint32_t delay = CalculateReconnectDelay();
    
    if (m_ReconnectTimer) {
        m_ReconnectTimer->expires_after(std::chrono::milliseconds(delay));
        m_ReconnectTimer->async_wait(asio::bind_executor(m_Strand, [this](const asio::error_code& ec) {
            if (!ec) {
                // 尝试重新连接
                asio::error_code socketEc;
                if (m_Socket.is_open()) {
                    m_Socket.close(socketEc);
                }
                
                DoConnect(m_Host, m_Port);
            }
        }));
    }
}

void TcpChannel::UpdateConnectionState(ConnectionState newState, const std::string& reason) {
    /* 功能：更新连接状态 - 使用基类实现
       参数：newState-新状态，reason-状态变更原因 */
    BaseNetworkChannel::UpdateConnectionState(newState, reason);
}

void TcpChannel::UpdateNetworkStats(uint64_t bytesSent, uint64_t bytesReceived, bool messageSent, bool messageReceived) {
    /* 功能：更新网络统计信息
       参数：bytesSent-发送字节数，bytesReceived-接收字节数
             messageSent-是否为消息发送，messageReceived-是否为消息接收 */
    std::lock_guard<std::mutex> lock(m_StatsMutex);
    
    m_Stats.TotalBytesSent += bytesSent;
    m_Stats.TotalBytesReceived += bytesReceived;
    
    if (messageSent) {
        m_Stats.MessagesSent++;
    }
    
    if (messageReceived) {
        m_Stats.MessagesReceived++;
    }
}

void TcpChannel::UpdateRtt(uint32_t rtt) {
    /* 功能：更新RTT统计
       参数：rtt-往返时间（毫秒） */
    std::lock_guard<std::mutex> lock(m_StatsMutex);
    
    m_RttSamples.push_back(rtt);
    
    // 保留最近10个RTT样本
    if (m_RttSamples.size() > 10) {
        m_RttSamples.erase(m_RttSamples.begin());
    }
    
    // 计算平均RTT
    if (!m_RttSamples.empty()) {
        uint64_t sum = 0;
        for (auto sample : m_RttSamples) {
            sum += sample;
        }
        m_Stats.AverageRtt = static_cast<uint32_t>(sum / m_RttSamples.size());
    }
    
    m_Stats.LatestRtt = rtt;
}

uint32_t TcpChannel::CalculateReconnectDelay() const {
    /* 功能：计算重连延迟（指数退避）
       返回：下次重连延迟（毫秒） */
    uint32_t delay = m_ReconnectInitialDelay;
    
    // 指数退避
    for (uint32_t i = 1; i < m_ReconnectAttempts && delay < m_ReconnectMaxDelay / 2; ++i) {
        delay *= 2;
    }
    
    return std::min(delay, m_ReconnectMaxDelay);
}

} // namespace SeanNetwork
